import math
import numpy as np

EPISODE_NUMBER = 10
STEP_NUMBER = 3000
beacon_transmit_power = 2  # 单位：W
beacon_num = 8  # 信标数
boundary = 1000  # 海域范围

def Average_acoustic_speed(depth_beacon, depth_target):
    #speed = 1480 - 0.1 * depth
    a = 0.1
    b = 1480
    speed = a * (depth_beacon - depth_target) / (math.log(b - a * depth_target) - math.log(b - a * depth_beacon))
    return speed


STATES = {
    'Depth': np.arange(0, 1000, 1000),
    'beacons': np.arange(0, 8, 1),
    'Localization_Error': np.arange(0, 1, 0.1),
}
count = 1
for k, v in STATES.items():
    count *= len(v)
STATE_NUMBER = count


# 0表示不选该信标，1表示选该信标
ACTIONS = {
    'Select_beacons_1': [0, 1],
    'Select_beacons_2': [0, 1],
    'Select_beacons_3': [0, 1],
    'Select_beacons_4': [0, 1],
    'Select_beacons_5': [0, 1],
    'Select_beacons_6': [0, 1],
    'Select_beacons_7': [0, 1],
    'Select_beacons_8': [0, 1],
}
count = 1
for k, v in ACTIONS.items():
    count *= len(v)
ACTION_NUMBER = count

def state_encode(state):
    Depth = state['Depth']
    beacons = state['beacons']
    Localization_Error = state['Localization_Error']
    Depth_idx = np.abs(Depth - STATES['Depth']).argmin()
    beacons_idx = np.abs(beacons - STATES['beacons']).argmin()
    Localization_Error_idx = np.abs(Localization_Error - STATES['Localization_Error']).argmin()
    state_idx = \
        Depth_idx + \
        beacons_idx * len(STATES['Depth']) + \
        Localization_Error_idx * len(STATES['Depth']) * len(STATES['beacons'])
    return state_idx

def action_encode(action):
    select_num = np.zeros(beacon_num)
    for i in range(beacon_num):
        i += 1
        select_num[i - 1] = action % 2
        action = (action - select_num[i-1]) / 2
    return select_num

def reward_function(r, p):
    u = - 2 * r - p
    return u

def reset():
    beacons_id = np.random.randint(len(STATES['beacons']))
    Localization_Error_id = np.random.randint(len(STATES['Localization_Error']))
    #Depth = np.random.randint(0, 1000)
    Depth = 600
    beacons = STATES['beacons'][beacons_id]
    Localization_Error = STATES['Localization_Error'][Localization_Error_id]
    results = {
        'Depth': Depth,
        'beacons': beacons,
        'Localization_Error': Localization_Error,
    }
    return results